#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Mon Feb 17 21:12:15 2020

@author: mario
"""
import sys
import time
import copy
import datetime
import numpy as np
import pandas as pd
import tushare as ts
from pathlib import Path
from collections import Counter


# In-history:
def get_history_price_range(DATA):
    # return `max` and `min` prices
    return DATA.price.max(), DATA.price.min()


# Grid-distribution analysis: revenue decays while the difference might increase.
def get_history_grid_size(DATA, gaps=np.arange(4, 11) / 1000):
    # :return: distribution of gaps in a specific fund product.
    def count_diff(seq):
        if isinstance(seq, list):
            seq = np.array(seq)
        _diffs = []
        for index, each in enumerate(seq):
            _diffs += [round(each, 3) for each in list(seq[index:] - each)]
        count = dict(Counter(_diffs))
        _count = []
        for key, value in count.items():
            _count += [[key, value]]
        return _count

    _stat = pd.DataFrame()
    for date in DATA.date.unique():
        _sub = DATA[DATA.date == date]
        # all prices in a day:
        _prices = _sub.price.unique()
        _count = count_diff(_prices)
        _count = pd.DataFrame(_count)
        _count.columns = ["diff_", "count_"]
        # daily return:
        _day_change = (_sub.price.values[-1] - _sub.price.values[0]) / _sub.price.values[0]
        _day_return = round(_day_change, 4)
        # stat: with three types
        _groupCount = {}
        _groupCount["date"] = date
        _groupCount["return"] = _day_return
        for gap in gaps:
            _field = "count_%d" % (gap * 1000)
            _groupCount[_field] = _count[_count.diff_ < gap].count_.sum()
        # table
        _groupCount = pd.DataFrame(_groupCount, index=[0])
        _stat = pd.concat([_stat, _groupCount], axis=0)
    # end
    _stat.index = range(len(_stat))
    _fields = ["count_%d" % (gap * 1000) for gap in gaps]
    return _stat[_fields].sum(axis=0).diff().dropna().to_dict()


def download_data(code):
    long = pd.DataFrame()
    _startday = "2018-01-01"
    day = copy.copy(_startday)
    _today = datetime.datetime.now().strftime("%Y-%m-%d")

    while day < _today:
        try:
            _stamp = datetime.datetime.strptime(day, "%Y-%m-%d")
            if (datetime.datetime.weekday(_stamp) > 0) and (datetime.datetime.weekday(_stamp) < 6):
                _tick = ts.get_tick_data(str(code), date=day, src='tt', retry_count=3)
                if isinstance(_tick, pd.DataFrame):
                    _tick["date"] = day
                    long = pd.concat([long, _tick], axis=0)
                    # 日期+1
                    sys.stdout.write("\r[Download Success] process at: [%s]" % day)
                    time.sleep(1.5)

                _stamp += datetime.timedelta(days=1)
                day = datetime.datetime.strftime(_stamp, "%Y-%m-%d")
                # 记录:
                sys.stdout.write("\r[Download Skip] process at: [%s]" % day)
                time.sleep(.5)

            else:
                _stamp += datetime.timedelta(days=1)
                day = datetime.datetime.strftime(_stamp, "%Y-%m-%d")
                next
        except Exception as e:
            print("[Download] %s" % e)
            break

    filename = (Path() / ("tushare_" + str(code))).with_suffix(".csv")
    long.to_csv(str(filename))
    print("[Download] stop at: [%s]" % day)
    sys.exit()


if __name__ == "__main__":
    # NB. the list of funds:
    FUNDS = dict(环保ETF=512580, 创业板50=159949, 券商ETF=510900,
                 H股ETF=510900, 芯片ETF=159995)
    name = "芯片ETF"
    print("[Download] start with %s: %d" % (name, FUNDS[name]))
    download_data(FUNDS[name])
